import importlib.abc
import importlib.util

from zkl_pyutils_fsspec import Fs

dataset_script_file_name = "dataset.py"
dataset_script_func_name = "load"


class DatasetScriptLoader(importlib.abc.InspectLoader):
    def __init__(self, source: str):
        self._source = source

    def get_source(self, fullname):
        return self._source


def load_dataset_script(fs: Fs):
    with fs.open(dataset_script_file_name, "rt") as fp:
        return fp.read()


def load_dataset_from_script(fs: Fs, *args, **kwargs):
    # load script
    source = load_dataset_script(fs)

    # exec script
    loader = DatasetScriptLoader(source)
    spec = importlib.util.spec_from_loader("__dataset__", loader)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)

    # exec func
    load_func = getattr(module, dataset_script_func_name)
    return load_func(fs, *args, **kwargs)
